import numpy as np
from sklearn.neighbors import NearestNeighbors
import copy
import random
from skimage import io
from graph import Graph, a_star
from parameter import *
from copy import deepcopy
import os

class Graph_generator:
    def __init__(self, map_size, k_size, plot=False, test=False):
        self.k_size = k_size
        self.graph = Graph()
        self.node_coords = None
        self.plot = plot
        self.x = []
        self.y = []
        self.map_x = map_size[1]
        self.map_y = map_size[0]
        self.dist_threshold = (((self.map_x-1)/(NUMX-1))**2 + ((self.map_y-1)/(NUMY-1))**2)**0.5 * 1.1
        self.uniform_points = self.generate_uniform_points()
        self.nodes_list = []
        self.reference_policy = None
        self.adjacent_matrix = []
        self.network_adjacent_matrix = []
        self.test = test

    def generate_graph(self, robot_belief, file_path):
        free_area = self.free_area(robot_belief)
        free_area_to_check = free_area[:, 0] + free_area[:, 1] * 1j
        uniform_points_to_check = self.uniform_points[:, 0] + self.uniform_points[:, 1] * 1j
        _, _, candidate_indices = np.intersect1d(free_area_to_check, uniform_points_to_check, return_indices=True)
        node_coords = self.uniform_points[candidate_indices]
        
        self.node_coords = self.unique_coords(node_coords).reshape(-1, 2)

        # generate the collision free graph
        self.find_k_neighbor_all_nodes(self.node_coords, robot_belief)

        ## generate reference policy
        graph = list(self.graph.edges.values())
        edge_inputs = []
        for i, node in enumerate(graph):
            node_edges = list(map(int, node))
            edge_inputs.append(node_edges)
            assert i == node_edges[0]
        
        ## generate adjacent matrix
        self.adjacent_matrix = self.calculate_edge_mask(edge_inputs)

        dist = [[9999 for i in range(len(self.adjacent_matrix))] for j in range(len(self.adjacent_matrix))]
        for i in range(len(self.adjacent_matrix)):
            for j in range(len(self.adjacent_matrix)):
                if i == j:
                    dist[i][j] = 0
                elif self.adjacent_matrix[i][j] == 0:
                    dist[i][j] = 1

        for k in range(len(self.adjacent_matrix)):
            for i in range(len(self.adjacent_matrix)):
                for j in range(len(self.adjacent_matrix)):
                    if dist[i][k] + dist[k][j] < dist[i][j]:
                        dist[i][j] = dist[i][k] + dist[k][j]

        self.network_adjacent_matrix = np.array(dist)
        print('adj_matrix: ', len(self.adjacent_matrix))
        if self.test:
            adj_file = 'adj_file/test_big'
        else:
            adj_file = 'adj_file/train_big'
        adjacent_matrix_file = str(adj_file + '/adjacent_matrix_'+file_path+'.txt')
        with open(adjacent_matrix_file, 'w') as f:
            f.write(str(len(self.adjacent_matrix))+'\n')
            for i in range(len(self.adjacent_matrix)):
                for j in range(len(self.adjacent_matrix)):
                    if j == len(self.adjacent_matrix)-1:
                        f.write(str(int(1-self.adjacent_matrix[i][j])) +'\n')
                    else:
                        f.write(str(int(1-self.adjacent_matrix[i][j])) +' ')

    def free_area(self, robot_belief):
            index = np.where(robot_belief == 255)
            free = np.asarray([index[1], index[0]]).T
            return free
    
    def unique_coords(self, coords):
        x = coords[:, 0] + coords[:, 1] * 1j
        indices = np.unique(x, return_index=True)[1]
        coords = np.array([coords[idx] for idx in sorted(indices)])
        return coords

    def find_k_neighbor_all_nodes(self, node_coords, robot_belief):
        X = node_coords
        if len(node_coords) >= self.k_size:
            knn = NearestNeighbors(n_neighbors=self.k_size)
        else:
            knn = NearestNeighbors(n_neighbors=len(node_coords))
        knn.fit(X)
        distances, indices = knn.kneighbors(X)

        for i, p in enumerate(X):
            # for j, neighbour in enumerate(X):
            for j, neighbour in enumerate(X[indices[i][:]]):
                start = p
                end = neighbour
                if (not self.check_collision(start, end, robot_belief)) or (not self.check_collision(end, start, robot_belief)):
                    a = str(self.find_index_from_coords(node_coords, p))
                    b = str(self.find_index_from_coords(node_coords, neighbour))
                    # dist = np.linalg.norm(p - neighbour)
                    # if (int(a) == 41 and int(b) == 35):
                    #     print('dis', dist)
                    if distances[i, j] < self.dist_threshold:
                        self.graph.add_node(a)
                        self.graph.add_edge(a, b, distances[i, j])
                        
    def check_collision(self, start, end, robot_belief):
        # Bresenham line algorithm checking
        collision = False

        x0 = start[0].round()
        y0 = start[1].round()
        x1 = end[0].round()
        y1 = end[1].round()
        dx, dy = abs(x1 - x0), abs(y1 - y0)
        x, y = x0, y0
        error = dx - dy
        x_inc = 1 if x1 > x0 else -1
        y_inc = 1 if y1 > y0 else -1
        dx *= 2
        dy *= 2

        while 0 <= x < robot_belief.shape[1] and 0 <= y < robot_belief.shape[0]:
            k = robot_belief.item(int(y), int(x))
            if x == x1 and y == y1:
                break
            if k == 1:
                collision = True
                break
            if k == 127:
                collision = True
                break
            if error > 0:
                x += x_inc
                error -= dy
            else:
                y += y_inc
                error += dx
        return collision
    
    def find_index_from_coords(self, node_coords, p):
        return np.where(np.linalg.norm(node_coords - p, axis=1) < 1e-5)[0][0]
    
    def calculate_edge_mask(self, edge_inputs):
        size = len(edge_inputs)
        bias_matrix = np.ones((size, size))
        for i in range(size):
            cnt = 0
            for j in range(size):
                if j in edge_inputs[i]:
                    bias_matrix[i][j] = 0
                    cnt += 1
            assert cnt <= K_SIZE
        return bias_matrix

    def generate_uniform_points(self):
        x = np.linspace(0, self.map_x - 1, NUMX).round().astype(int) # 55
        y = np.linspace(0, self.map_y - 1, NUMY).round().astype(int)
        t1, t2 = np.meshgrid(x, y)
        points = np.vstack([t1.T.ravel(), t2.T.ravel()]).T
        return points
    
class Env():
    def __init__(self, map_index, num_robots, random_seed=None, k_size=20, plot=False, test=False):
        # import environment ground truth from dungeon files
        self.test = test
        if self.test:
            self.map_dir = f'DungeonMaps/test_1025'  # change to 'complex', 'medium', and 'easy'
        else:
            self.map_dir = f'DungeonMaps/train_1025'
        self.random_seed = random_seed
        self.map_list = os.listdir(self.map_dir)
        self.map_list.sort(reverse=False)
        self.map_index = map_index % np.size(self.map_list)
        self.ground_truth, self.start_position, self.target_position = self.import_ground_truth(
            self.map_dir + '/' + self.map_list[self.map_index])

        self.ground_truth_size = np.shape(self.ground_truth)  # (480, 640)
        self.num_robots = num_robots
        self.robot_positions = []
        self.robot_indexes = []
        for k in range(num_robots):
            self.robot_positions.append(self.start_position)
        self.target_index = -1
        # initialize graph generator
        self.graph_generator = Graph_generator(map_size=self.ground_truth_size, k_size=k_size, plot=plot, test=self.test)
        if FIXED_OPPONENT == False:
            self.node_coords, self.graph, self.node_feature, self.reference_policy, self.adjacent_matrix, self.network_adjacent_matrix, self.opponent_policy = None, None, None, None, None, None, None
        else:
            self.node_coords, self.graph, self.node_feature, self.reference_policy, self.adjacent_matrix, self.network_adjacent_matrix = None, None, None, None, None, None
        self.stepi = 0
        self.node_num = 0

        self.begin()
       
    def find_index_from_coords(self, position):
        index = np.argmin(np.linalg.norm(self.node_coords - position, axis=1))
        assert self.node_num == np.linalg.norm(self.node_coords - position, axis=1).shape[0]
        return index

    def begin(self):
        file_path = str(self.map_list[self.map_index])

        self.graph_generator.generate_graph(self.ground_truth, file_path)
        
    def import_ground_truth(self, map_index):
        # occupied 1, free 255, unexplored 127
        ground_truth = (io.imread(map_index, 1) * 255).astype(int)
        if self.random_seed != None: 
            robot_location = np.nonzero(ground_truth == 208)
            robot_location = np.array([np.array(robot_location)[1, 127], np.array(robot_location)[0, 127]])
            
            target_position = np.nonzero(ground_truth == 194)
            random.seed(self.random_seed)
            random_number = random.randint(1, np.array(target_position).shape[1] - 1)
            target_position = np.array([np.array(target_position)[1, random_number], np.array(target_position)[0, random_number]])
           
        else: 
            count = 0
            threshhold = (((np.shape(ground_truth)[1]-1)/(NUMX-1))**2 + ((np.shape(ground_truth)[0]-1)/(NUMY-1))**2)**0.5 * 1.1
            while True:
                robot_indices = np.nonzero(ground_truth == 208)
                target_indices = np.nonzero(ground_truth == 194)
                free_area = np.concatenate((np.array(robot_indices), np.array(target_indices)), axis=1)

                random_number_1 = random.randint(1, np.array(free_area).shape[1] - 1)
                random_number_2 = random.randint(1, np.array(free_area).shape[1] - 1)
                robot_location = np.array([np.array(free_area)[1, random_number_1], np.array(free_area)[0, random_number_1]])
                target_position = np.array([np.array(free_area)[1, random_number_2], np.array(free_area)[0, random_number_2]])
            
                dist = np.linalg.norm(robot_location - target_position)
                count += 1
                if dist > threshhold * 5 or count > 1000:
                    break
        
        ground_truth = (ground_truth > 150)
        # save_ground_truth = ground_truth.copy()
        # io.imsave(map_index.split('/')[-1][:-4]+'_clean.png', save_ground_truth)
        ground_truth = ground_truth * 254 + 1

        return ground_truth, robot_location, target_position

if __name__=='__main__':
    for i in range(76):
        env = Env(map_index=i, k_size=K_SIZE, num_robots=N_ROBOTS, random_seed=RANDOM_SEED, test=False)